﻿#include "StaticWebServer.h"

#include <fstream>
#include <sstream>
#include <thread>
#include <algorithm>
#include <filesystem>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#endif

NahidaProject::StaticWebServer::StaticWebServer(int port, const std::string& rootDirectory) : port(port), rootDirectory(rootDirectory), running(false) {
#ifdef _WIN32
	WSADATA wsaData;
	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
		throw std::runtime_error("Failed to initialize Winsock");
	}
#endif

	// 创建socket
	serverSocket = socket(AF_INET, SOCK_STREAM, 0);
	if (serverSocket < 0) {
		Cleanup();
		throw std::runtime_error("Failed to create socket");
	}

	// 设置socket选项，允许地址重用
	int opt = 1;
	if (setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR,
		reinterpret_cast<const char*>(&opt), sizeof(opt)) < 0) {
		Cleanup();
		throw std::runtime_error("Failed to set socket options");
	}

	// 绑定地址
	sockaddr_in serverAddress{};
	serverAddress.sin_family = AF_INET;
	serverAddress.sin_addr.s_addr = INADDR_ANY;
	serverAddress.sin_port = htons(port);

	if (bind(serverSocket, reinterpret_cast<sockaddr*>(&serverAddress),
		sizeof(serverAddress)) < 0) {
		Cleanup();
		throw std::runtime_error("Failed to bind socket");
	}

	// 监听连接
	if (listen(serverSocket, 10) < 0) {
		Cleanup();
		throw std::runtime_error("Failed to listen on socket");
	}

	// 确保根目录存在
	if (!std::filesystem::exists(rootDirectory)) {
		std::filesystem::create_directories(rootDirectory);
	}

	std::cout << "Server initialized successfully" << std::endl;
}

NahidaProject::StaticWebServer::~StaticWebServer() {
	StopServer();
	Cleanup();
}

void NahidaProject::StaticWebServer::StartServer() {
	running = true;

	while (running) {
		sockaddr_in clientAddress{};
		socklen_t clientLength = sizeof(clientAddress);

		// 接受连接
		int clientSocket = accept(serverSocket, reinterpret_cast<sockaddr*>(&clientAddress), &clientLength);

		if (clientSocket < 0) {
			if (running) {
				std::cerr << "Failed to accept connection" << std::endl;
			}
			continue;
		}

		// 为每个连接创建新线程处理
		std::thread([this, clientSocket]() {
			HandleClient(clientSocket);
		}).detach();
	}
}

void NahidaProject::StaticWebServer::StopServer() {
	running = false;

	// 关闭监听socket以中断accept()
	if (serverSocket >= 0) {
#ifdef _WIN32
		closesocket(serverSocket);
#else
		close(serverSocket);
#endif
		serverSocket = -1;
	}
}

void NahidaProject::StaticWebServer::HandleClient(int clientSocket) {
	try {
		// 读取请求
		std::string request = ReadRequest(clientSocket);
		if (request.empty()) {
			SendErrorResponse(clientSocket, 400, "Bad Request");
			return;
		}

		// 解析请求
		StaticWebServerHTTPRequest req = ParseRequest(request);

		// 处理请求
		if (req.method == "GET" || req.method == "HEAD") {
			HandleGetRequest(clientSocket, req);
		}
		else {
			SendErrorResponse(clientSocket, 405, "Method Not Allowed");
		}

	}
	catch (const std::exception& e) {
		std::cerr << "Error handling client: " << e.what() << std::endl;
		SendErrorResponse(clientSocket, 500, "Internal Server Error");
	}

	// 关闭客户端连接
#ifdef _WIN32
	closesocket(clientSocket);
#else
	close(clientSocket);
#endif
}

std::string NahidaProject::StaticWebServer::ReadRequest(int clientSocket) {
	std::string request;
	char buffer[1024];
	int bytesReceived;

	// 设置非阻塞模式并读取数据
	while ((bytesReceived = recv(clientSocket, buffer, sizeof(buffer) - 1, 0)) > 0) {
		buffer[bytesReceived] = '\0';
		request += buffer;

		// 检查是否收到完整的HTTP头部
		if (request.find("\r\n\r\n") != std::string::npos) {
			break;
		}
	}

	return request;
}

NahidaProject::StaticWebServer::StaticWebServerHTTPRequest NahidaProject::StaticWebServer::ParseRequest(const std::string& request) {
	StaticWebServerHTTPRequest req;

	std::istringstream iss(request);
	std::string line;

	// 解析第一行：方法 URL 协议版本
	if (std::getline(iss, line)) {
		std::istringstream lineStream(line);
		lineStream >> req.method >> req.url >> req.version;
	}

	// 解析头部字段
	while (std::getline(iss, line) && line != "\r") {
		size_t colonPosition = line.find(':');
		if (colonPosition != std::string::npos) {
			std::string key = line.substr(0, colonPosition);
			std::string value = line.substr(colonPosition + 2); // 跳过冒号和空格

			// 移除回车符
			if (!value.empty() && value.back() == '\r') {
				value.pop_back();
			}

			req.headers[key] = value;
		}
	}

	return req;
}

void NahidaProject::StaticWebServer::HandleGetRequest(int clientSocket, const StaticWebServerHTTPRequest& request) {
	std::string path = request.url;

	// 移除查询字符串
	size_t queryPosition = path.find('?');
	if (queryPosition != std::string::npos) {
		path = path.substr(0, queryPosition);
	}

	// 处理默认页面
	if (path == "/" || path.empty()) {
		path = "/index.html";
	}

	// 构建文件路径
	std::string filePath = rootDirectory + path;

	// 规范化路径，防止目录遍历攻击
	std::filesystem::path normalizedPath(filePath);
	std::filesystem::path rootPath(rootDirectory);

	if (normalizedPath.string().find(rootPath.string()) != 0) {
		SendErrorResponse(clientSocket, 403, "Forbidden");
		return;
	}

	// 检查文件是否存在
	if (!std::filesystem::exists(normalizedPath)) {
		SendErrorResponse(clientSocket, 404, "Not Found");
		return;
	}

	// 如果是目录，尝试查找index文件
	if (std::filesystem::is_directory(normalizedPath)) {
		std::vector<std::string> indexFiles = { "index.html", "index.htm" };
		bool found = false;

		for (const auto& indexFile : indexFiles) {
			std::filesystem::path indexPath = normalizedPath / indexFile;
			if (std::filesystem::exists(indexPath)) {
				normalizedPath = indexPath;
				found = true;
				break;
			}
		}

		if (!found) {
			SendDirectoryListing(clientSocket, normalizedPath);
			return;
		}
	}

	SendFile(clientSocket, normalizedPath, request.method == "HEAD");
}

void NahidaProject::StaticWebServer::SendFile(int clientSocket, const std::filesystem::path& filePath, bool headOnly) {
	std::ifstream file(filePath, std::ios::binary);
	if (!file.is_open()) {
		SendErrorResponse(clientSocket, 500, "Internal Server Error");
		return;
	}

	// 获取文件大小
	file.seekg(0, std::ios::end);
	size_t fileSize = file.tellg();
	file.seekg(0, std::ios::beg);

	// 获取MIME类型
	std::string mimeType = GetMimeType(filePath.extension().string());

	// 构造响应头
	std::ostringstream responseHeader;
	responseHeader << "HTTP/1.1 200 OK\r\n";
	responseHeader << "Content-Type: " << mimeType << "\r\n";
	responseHeader << "Content-Length: " << fileSize << "\r\n";
	responseHeader << "Connection: close\r\n";
	responseHeader << "\r\n";

	// 发送响应头
	std::string headerString = responseHeader.str();
	send(clientSocket, headerString.c_str(), headerString.length(), 0);

	// 如果是HEAD请求，只发送头部
	if (headOnly) {
		return;
	}

	// 发送文件内容
	const size_t bufferSize = 8192;
	char buffer[bufferSize];

	while (file.read(buffer, bufferSize) || file.gcount() > 0) {
		send(clientSocket, buffer, file.gcount(), 0);
	}
}

void NahidaProject::StaticWebServer::SendDirectoryListing(int clientSocket, const std::filesystem::path& directoryPath) {
	std::ostringstream html;
	html << "<!DOCTYPE html>\n<html>\n<head><title>Directory Listing</title></head>\n<body>\n";
	html << "<h1>Directory Listing</h1>\n<ul>\n";

	// 添加父目录链接
	html << "<li><a href=\"../\">../</a></li>\n";

	// 列出目录内容
	for (const auto& entry : std::filesystem::directory_iterator(directoryPath)) {
		std::string name = entry.path().filename().string();
		std::string link = name + (entry.is_directory() ? "/" : "");
		html << "<li><a href=\"" << link << "\">" << link << "</a></li>\n";
	}

	html << "</ul>\n</body>\n</html>";

	std::string content = html.str();

	// 构造响应
	std::ostringstream responseHeader;
	responseHeader << "HTTP/1.1 200 OK\r\n";
	responseHeader << "Content-Type: text/html\r\n";
	responseHeader << "Content-Length: " << content.length() << "\r\n";
	responseHeader << "Connection: close\r\n";
	responseHeader << "\r\n";

	std::string response = responseHeader.str() + content;
	send(clientSocket, response.c_str(), response.length(), 0);
}

void NahidaProject::StaticWebServer::SendErrorResponse(int clientSocket, int statusCode, const std::string& message) {
	std::string statusText;
	switch (statusCode) {
	case 400: statusText = "Bad Request"; break;
	case 403: statusText = "Forbidden"; break;
	case 404: statusText = "Not Found"; break;
	case 405: statusText = "Method Not Allowed"; break;
	case 500: statusText = "Internal Server Error"; break;
	default: statusText = "Unknown Error"; break;
	}

	std::ostringstream html;
	html << "<!DOCTYPE html>\n<html>\n<head><title>" << statusCode << " " << statusText << "</title></head>\n";
	html << "<body>\n<h1>" << statusCode << " " << statusText << "</h1>\n";
	html << "<p>" << message << "</p>\n</body>\n</html>";

	std::string content = html.str();

	std::ostringstream responseHeader;
	responseHeader << "HTTP/1.1 " << statusCode << " " << statusText << "\r\n";
	responseHeader << "Content-Type: text/html\r\n";
	responseHeader << "Content-Length: " << content.length() << "\r\n";
	responseHeader << "Connection: close\r\n";
	responseHeader << "\r\n";

	std::string response = responseHeader.str() + content;
	send(clientSocket, response.c_str(), response.length(), 0);
}

std::string NahidaProject::StaticWebServer::GetMimeType(const std::string& extension) {
	static const std::map<std::string, std::string> mimeTypes = {
		{".html", "text/html"},
		{".htm", "text/html"},
		{".css", "text/css"},
		{".js", "application/javascript"},
		{".json", "application/json"},
		{".xml", "application/xml"},
		{".txt", "text/plain"},
		{".png", "image/png"},
		{".jpg", "image/jpeg"},
		{".jpeg", "image/jpeg"},
		{".gif", "image/gif"},
		{".ico", "image/x-icon"},
		{".svg", "image/svg+xml"},
		{".pdf", "application/pdf"},
		{".zip", "application/zip"},
		{".mp3", "audio/mpeg"},
		{".mp4", "video/mp4"}
	};

	auto it = mimeTypes.find(extension);
	if (it != mimeTypes.end()) {
		return it->second;
	}

	return "application/octet-stream";
}

void NahidaProject::StaticWebServer::Cleanup() {
#ifdef _WIN32
	WSACleanup();
#endif
}
